#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
语音传感器模块

该模块负责音频采集、语音活动检测(VAD)和唤醒词检测。
使用webrtcvad进行语音活动检测，使用pvporcupine进行唤醒词检测，使用pyaudio进行音频采集。
"""

import os
import time
import numpy as np
import webrtcvad
import pyaudio
import wave
import threading
from pathlib import Path
import struct
import pvporcupine
import alsaaudio
from typing import Optional, Callable, Union, Dict, Any, List
from dataclasses import dataclass
from enum import Enum


class VoiceEventType(Enum):
    """语音事件类型"""

    WAKE_WORD = "wake_word"  # 唤醒词检测
    SPEECH_START = "speech_start"  # 语音开始
    SPEECH_DATA = "speech_data"  # 语音数据
    SPEECH_END = "speech_end"  # 语音结束
    ERROR = "error"  # 错误事件


@dataclass
class VoiceEvent:
    """语音事件数据类"""

    type: VoiceEventType
    data: Optional[bytes] = None
    timestamp: float = 0.0
    error: Optional[str] = None


class VoiceSensor:
    """
    语音传感器类

    负责音频采集、语音活动检测和唤醒词检测
    """

    def __init__(self, config: Dict[str, Any], callback: Optional[Callable] = None):
        """
        初始化语音传感器

        Args:
            config: 配置字典，包含以下参数：
                - access_key: Porcupine访问密钥
                - wake_word: 唤醒词
                - vad_aggressiveness: VAD灵敏度(0-3)
                - sample_rate: 采样率
                - frame_duration_ms: 音频帧长度(毫秒)
                - input_device: 音频输入设备
                - max_silence_seconds: 最大静音时间(秒)
                - min_speech_seconds: 最小语音时间(秒)
                - save_audio: 是否保存音频文件
                - audio_dir: 音频文件保存目录
            callback: 回调函数，用于处理语音事件
        """
        self.config = config
        self.callback = callback

        # 从配置中获取参数
        self.access_key = config.get("access_key")
        self.wake_word = config.get("wake_word", "hey siri")
        self.vad_aggressiveness = config.get("vad_aggressiveness", 2)
        self.sample_rate = config.get("sample_rate", 16000)
        self.frame_duration_ms = config.get("frame_duration_ms", 30)
        self.frame_size = int(self.sample_rate * self.frame_duration_ms / 1000)
        self.input_device = config.get("input_device", "plughw:2,0")
        self.max_silence_seconds = config.get("max_silence_seconds", 20)
        self.min_speech_seconds = config.get("min_speech_seconds", 0.5)
        self.save_audio = config.get("save_audio", False)
        self.audio_dir = config.get("audio_dir", "audio")

        # 确保frame_size是10ms的倍数
        if self.frame_size % 160 != 0:
            raise ValueError("frame_size必须是10ms的倍数")

        # 状态变量
        self.is_listening = False
        self.is_recording = False
        self.audio_buffer = []
        self.buffer_max_size = 50
        self.speech_frames = 0
        self.silence_frames = 0
        self.min_speech_frames = int(
            self.min_speech_seconds * 1000 / self.frame_duration_ms
        )
        self.min_silence_frames = int(0.5 * 1000 / self.frame_duration_ms)  # 0.5秒静音
        self.event_history: List[VoiceEvent] = []
        self.max_history_size = 100

        # 初始化组件
        self._init_vad()
        self._init_porcupine()
        self._init_audio_input()

        # 创建音频保存目录
        if self.save_audio:
            os.makedirs(self.audio_dir, exist_ok=True)

        print(f"语音传感器初始化完成，VAD灵敏度: {self.vad_aggressiveness}")

    def _init_vad(self):
        """初始化语音活动检测器"""
        self.vad = webrtcvad.Vad(self.vad_aggressiveness)

    def _init_porcupine(self):
        """初始化Porcupine唤醒词检测器"""
        self.porcupine = None
        if self.access_key:
            try:
                self.porcupine = pvporcupine.create(
                    access_key=self.access_key, keywords=[self.wake_word]
                )
                print(f"Porcupine初始化成功，唤醒词: {self.wake_word}")
            except Exception as e:
                print(f"Porcupine初始化失败: {e}")
                import traceback

                traceback.print_exc()
                print("唤醒词检测功能将不可用")
        else:
            print("未提供Porcupine访问密钥，唤醒词检测功能将不可用")

    def _init_audio_input(self):
        """初始化ALSA音频输入设备"""
        max_retries = 3
        retry_delay = 1.0

        for attempt in range(max_retries):
            try:
                if self.input is not None:
                    try:
                        self.input.close()
                    except:
                        pass
                    self.input = None
                    time.sleep(0.5)

                self.input = alsaaudio.PCM(
                    alsaaudio.PCM_CAPTURE,
                    alsaaudio.PCM_NORMAL,
                    device=self.input_device,
                )

                self.input.setchannels(1)
                self.input.setrate(self.sample_rate)
                self.input.setformat(alsaaudio.PCM_FORMAT_S16_LE)
                self.input.setperiodsize(160)

                print(f"音频输入设备初始化成功: {self.input_device}")
                return True

            except Exception as e:
                print(f"音频输入设备初始化失败 (尝试 {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    print(f"等待 {retry_delay} 秒后重试...")
                    time.sleep(retry_delay)
                else:
                    print("达到最大重试次数，初始化失败")
                    self.input = None
                    return False

    def is_speech(self, frame: bytes) -> bool:
        """
        检测音频帧是否包含语音

        Args:
            frame: 音频帧数据

        Returns:
            bool: 是否包含语音
        """
        try:
            frame_length = len(frame)
            if frame_length % 160 != 0:
                target_length = (frame_length // 160) * 160
                if frame_length < target_length:
                    frame = frame + b"\x00" * (target_length - frame_length)
                else:
                    frame = frame[:target_length]

            if len(frame) < 160 or len(frame) > 480:
                return False

            return self.vad.is_speech(frame, self.sample_rate)
        except Exception as e:
            print(f"VAD检测错误: {e}")
            return False

    def detect_wake_word(self, frame: bytes) -> bool:
        """
        检测音频帧中是否包含唤醒词

        Args:
            frame: 音频帧数据

        Returns:
            bool: 是否检测到唤醒词
        """
        if self.porcupine is None:
            return False

        try:
            pcm = struct.unpack_from("h" * (len(frame) // 2), frame)
            keyword_index = self.porcupine.process(pcm)

            if keyword_index >= 0:
                print(f"检测到唤醒词: {self.wake_word}")
                return True

            return False
        except Exception as e:
            print(f"唤醒词检测错误: {e}")
            return False

    def _add_event(self, event: VoiceEvent):
        """添加事件到历史记录"""
        event.timestamp = time.time()
        self.event_history.append(event)
        if len(self.event_history) > self.max_history_size:
            self.event_history.pop(0)

    def _notify_callback(self, event: VoiceEvent):
        """通知回调函数"""
        if self.callback:
            try:
                self.callback(event)
            except Exception as e:
                print(f"执行回调时出错: {e}")
                self._add_event(VoiceEvent(VoiceEventType.ERROR, error=str(e)))

    def start(self):
        """开始监听音频"""
        if self.is_listening:
            print("已经在监听中")
            return

        if self.input is None:
            if not self._init_audio_input():
                print("无法初始化音频设备，无法开始监听")
                return

        self.is_listening = True
        self.audio_buffer = []
        self.speech_frames = 0
        self.silence_frames = 0

        print("开始监听语音活动和唤醒词...")

        self.listen_thread = threading.Thread(target=self._listen_loop)
        self.listen_thread.daemon = True
        self.listen_thread.start()

    def _listen_loop(self):
        """监听循环，处理音频数据"""
        try:
            while self.is_listening:
                try:
                    l, frame = self.input.read()
                    if l <= 0:
                        continue

                    if len(frame) != 160 * 2:
                        continue

                    if self.is_speech(frame):
                        self.speech_frames += 1
                        self.silence_frames = 0

                        if self.is_recording:
                            self.audio_buffer.append(frame)

                        if self.detect_wake_word(frame):
                            event = VoiceEvent(VoiceEventType.WAKE_WORD)
                            self._add_event(event)
                            self._notify_callback(event)

                        elif (
                            self.speech_frames >= self.min_speech_frames
                            and not self.is_recording
                        ):
                            print(f"检测到语音活动! 开始录音")
                            self.is_recording = True
                            self.audio_buffer = [frame]

                            event = VoiceEvent(VoiceEventType.SPEECH_START)
                            self._add_event(event)
                            self._notify_callback(event)
                    else:
                        self.silence_frames += 1
                        self.speech_frames = 0

                        if (
                            self.is_recording
                            and self.silence_frames >= self.min_silence_frames
                        ):
                            silence_seconds = (
                                self.silence_frames * self.frame_duration_ms / 1000
                            )
                            if silence_seconds >= self.max_silence_seconds:
                                print(
                                    f"检测到语音结束! 连续静音: {silence_seconds:.1f}秒"
                                )
                                self._process_speech_segment()
                                self.is_recording = False
                                self.audio_buffer = []

                except alsaaudio.ALSAAudioError as e:
                    print(f"ALSA音频错误: {e}")
                    if not self._init_audio_input():
                        print("无法重新初始化音频设备，停止监听")
                        break
                    time.sleep(0.1)

                except Exception as e:
                    print(f"监听循环中发生错误: {e}")
                    import traceback

                    traceback.print_exc()
                    time.sleep(0.1)

        except Exception as e:
            print(f"监听循环中发生严重错误: {e}")
            import traceback

            traceback.print_exc()
        finally:
            self.is_listening = False
            self.is_recording = False

    def _process_speech_segment(self):
        """处理完整的语音片段"""
        if not self.audio_buffer:
            return

        print(f"处理语音片段，帧数: {len(self.audio_buffer)}")
        audio_data = b"".join(self.audio_buffer)

        # 保存音频文件
        if self.save_audio:
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            filename = os.path.join(self.audio_dir, f"speech_{timestamp}.wav")
            self._save_audio_file(audio_data, filename)

        # 通知回调
        event = VoiceEvent(VoiceEventType.SPEECH_DATA, data=audio_data)
        self._add_event(event)
        self._notify_callback(event)

        # 发送语音结束事件
        end_event = VoiceEvent(VoiceEventType.SPEECH_END)
        self._add_event(end_event)
        self._notify_callback(end_event)

    def _save_audio_file(self, audio_data: bytes, filename: str):
        """保存音频文件"""
        try:
            with wave.open(filename, "wb") as wf:
                wf.setnchannels(1)
                wf.setsampwidth(2)  # 16-bit
                wf.setframerate(self.sample_rate)
                wf.writeframes(audio_data)
            print(f"音频文件已保存: {filename}")
        except Exception as e:
            print(f"保存音频文件失败: {e}")

    def stop(self):
        """停止监听"""
        self.is_listening = False
        self.is_recording = False

        if hasattr(self, "listen_thread") and self.listen_thread.is_alive():
            self.listen_thread.join(timeout=1.0)

        self.audio_buffer = []

        if self.input is not None:
            try:
                self.input.close()
            except:
                pass
            self.input = None

        print("停止监听语音活动和唤醒词")

    def get_event_history(self, limit: Optional[int] = None) -> List[VoiceEvent]:
        """
        获取事件历史记录

        Args:
            limit: 返回的最大事件数量

        Returns:
            List[VoiceEvent]: 事件历史记录
        """
        if limit is None:
            return self.event_history.copy()
        return self.event_history[-limit:].copy()

    def clear_event_history(self):
        """清除事件历史记录"""
        self.event_history.clear()

    def set_callback(self, callback: Callable):
        """
        设置回调函数

        Args:
            callback: 回调函数，用于处理语音事件
        """
        self.callback = callback

    def dispose(self):
        """释放资源"""
        self.stop()

        if self.porcupine is not None:
            self.porcupine.delete()

        if hasattr(self, "audio"):
            self.audio.terminate()

        print("语音传感器已释放")


# 测试代码
if __name__ == "__main__":

    def on_voice_event(event: VoiceEvent):
        if event.type == VoiceEventType.WAKE_WORD:
            print("检测到唤醒词! 执行唤醒操作...")
        elif event.type == VoiceEventType.SPEECH_START:
            print("检测到语音活动! 开始录音...")
        elif event.type == VoiceEventType.SPEECH_DATA:
            print(f"收到语音数据，长度: {len(event.data)} 字节")
        elif event.type == VoiceEventType.SPEECH_END:
            print("语音结束")
        elif event.type == VoiceEventType.ERROR:
            print(f"发生错误: {event.error}")
        print("执行后续操作...")

    # 配置参数
    config = {
        "access_key": os.environ.get("PICOVOICE_ACCESS_KEY"),
        "wake_word": "hey siri",
        "vad_aggressiveness": 2,
        "sample_rate": 16000,
        "frame_duration_ms": 30,
        "input_device": "plughw:2,0",
        "max_silence_seconds": 4,
        "min_speech_seconds": 0.5,
        "save_audio": True,
        "audio_dir": "audio",
    }

    # 初始化语音传感器
    voice_sensor = VoiceSensor(config)
    voice_sensor.set_callback(on_voice_event)

    try:
        voice_sensor.start()
        print("按Ctrl+C退出")
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        voice_sensor.stop()
        print("程序退出")
